import numpy as np
import matplotlib.pyplot as plt

from sklearn.preprocessing import LabelBinarizer, OneHotEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.layers import Dense
from sklearn import datasets


def main(result_save_on_disk: bool = True):
    print('[info]:正在加载mnist(full)数据集')
    raw_data = datasets.fetch_openml('mnist_784')
    data = raw_data.data.astype("float") / 255.0  # 归一化

    # 独热编码
    # hot = OneHotEncoder()
    # label = hot.fit_transform([[i] for i in raw_data.target]).toarray()

    train_x, test_x, train_y, test_y = train_test_split(data, raw_data.target,
                                                        test_size=0.25)  # label
    train_x = train_x.values
    test_x = test_x.values
    # 热编码过程
    label_b = LabelBinarizer()
    train_y = label_b.fit_transform(train_y)
    test_y = label_b.transform(test_y)

    # 模型建立
    model = Sequential()
    model.add(Dense(256, input_shape=(784,), activation="sigmoid"))
    model.add(Dense(128, activation="sigmoid"))
    model.add(Dense(10, activation="softmax"))
    # 开始训练
    print('[info]:开始训练网络..')
    sgd = SGD(0.01)
    model.compile(loss="categorical_crossentropy",
                  optimizer=sgd,
                  metrics=["accuracy"])
    record_data = model.fit(train_x, train_y, validation_data=(test_x, test_y),
                            epochs=100, batch_size=128)
    # 评价网络
    print('[info]:正在评估网络...')
    predictions = model.predict(test_x, batch_size=128)
    print(classification_report(test_y.argmax(axis=1),
                                predictions.argmax(axis=1),
                                target_names=[str(i) for i in label_b.classes_]))  # hot.categories_[0]
    # 画图(注意不同版本的keras参数不一样，acc和accuracy)
    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, 100), record_data.history["loss"], label="train_loss")
    plt.plot(np.arange(0, 100), record_data.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, 100), record_data.history["accuracy"], label="train_acc")
    plt.plot(np.arange(0, 100), record_data.history["val_accuracy"], label="val_acc")
    plt.title("train_pac loss and accuracy")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend()
    if result_save_on_disk:
        plt.savefig("NN网络训练损失及准确率的变化图")


if __name__ == '__main__':
    main()
